from rdkit import Chem
from rdkit.Chem import Descriptors, Draw
import matplotlib.pyplot as plt
import numpy as np
import logging
import json
from typing import Dict, Union
from mordred import Calculator, descriptors

def configure_logging() -> None:
    """配置日志记录系统"""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler("solubility.log"),
            logging.StreamHandler()
        ]
    )

def visualize_molecule(mol: Chem.Mol) -> None:
    """生成分子结构可视化"""
    try:
        img = Draw.MolToImage(mol, size=(400, 300), kekulize=True)
        plt.figure(figsize=(6,4))
        plt.imshow(img)
        plt.axis('off')
        plt.title("Molecular Structure", fontsize=14)
        plt.savefig("molecule.png", bbox_inches='tight')
        plt.close()
    except Exception as e:
        logging.error(f"分子可视化失败: {str(e)}")

def predict_water_solubility(smiles: str) -> Union[Dict, None]:
    """使用RDKit预测分子水溶性并提供分析"""
    try:
        # 创建分子对象并验证
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            logging.warning(f"无效的SMILES: {smiles}")
            return None
        
        # 计算核心性质
        # logs = Descriptors.MolLogS(mol)
        # 正确初始化计算器
        calc = Calculator([descriptors.LogS])
        logs = calc(mol)[0]
        
        # 水溶性分类
        # LogS值解读: LogS是水溶性的对数表示，数值越高，水溶性越好
        # LogS = -2表示溶解度约为0.01 mol/L (1% 摩尔浓度) 
        # 每降低1个单位，溶解度降低10倍 
        # 水溶性分类参考: 高水溶性 (LogS > -1): > 0.1 mol/L 
        # 良好水溶性 (-1 > LogS > -3): 0.001-0.1 mol/L 
        # 中等水溶性 (-3 > LogS > -5): 0.00001-0.001 mol/L 
        # 低水溶性 (LogS < -5): < 0.00001 mol/L
        solubility_class = ""
        solubility_desc = ""
        if logs > -1:
            solubility_class = "高水溶性"
            solubility_desc = "溶解性良好，适合水溶液配方"
        elif logs > -3:
            solubility_class = "良好水溶性"
            solubility_desc = "溶解性良好，普通配方即可使用"
        elif logs > -5:
            solubility_class = "中等水溶性"
            solubility_desc = "溶解性可接受，可能需要溶解促进剂"
        else:
            solubility_class = "低水溶性"
            solubility_desc = "溶解性差，需要特殊配方技术"
        
        # 分子量
        mw = Descriptors.MolWt(mol)

        # 计算衍生性质：溶解度
        # 溶解度单位换算: mol/L: 每升水中溶解的摩尔数 
        # mg/L: 每升水中溶解的毫克数 
        # 转换公式: mg/L = mol/L × 分子量 × 1000 
        solubility_mol_L = 10**logs
        solubility_mg_L = solubility_mol_L * mw * 1000
        
        # LogP: 脂水分配系数，值越高越亲脂，水溶性越差 
        logp = Descriptors.MolLogP(mol)
        logp_analysis = ""
        # LogP < 1: 水溶性好 
        # LogP 1-3: 水溶性中等 
        # LogP > 3: 水溶性较差 
        # LogP > 5: 水溶性极差 
        if logp > 5:
            logp_analysis = "LogP值过高，极大降低水溶性"
        elif logp > 3:
            logp_analysis = "LogP值较高，明显降低水溶性"
        elif logp > 1:
            logp_analysis = "LogP值中等，轻微降低水溶性"
        else:
            logp_analysis = "LogP值较低，有利于水溶性"
        
        mw_analysis = ""
        # 分子量(MW): MW < 300: 有利于水溶性 
        # MW 300-500: 中性影响 
        # MW > 500: 不利于水溶性 
        if mw > 500:
            mw_analysis = "分子量大，不利于水溶性"
        elif mw >= 300:
            mw_analysis = "分子量中等，对水溶性影响中性"
        else:
            mw_analysis = "分子量小，有利于水溶性"
        
        tpsa = Descriptors.TPSA(mol)
        tpsa_analysis = ""
        # 极性表面积(TPSA): TPSA < 40 Å²: 不利于水溶性 
        # TPSA 40-90 Å²: 有利于水溶性 
        # TPSA > 90 Å²: 极有利于水溶性
        if tpsa > 90:
            tpsa_analysis = "极性表面积大，极有利于水溶解"
        elif tpsa >= 40:
            tpsa_analysis = "极性表面积适中，有利于水溶解"
        else:
            tpsa_analysis = "极性表面积小，不利于水溶解"
        
        # 氢键供体/受体数: 数量越多，越有利于形成与水的氢键，水溶性越好 
        # 氢键供体数
        hbd = Descriptors.NumHDonors(mol)
        # 氢键受体数
        hba = Descriptors.NumHAcceptors(mol)

        result = {
            "smiles": smiles,
            "logS": round(logs, 3),
            "solubility": {
                "mol_L": "{:.6e}".format(solubility_mol_L),
                "mg_L": round(solubility_mg_L, 2),
                "class": solubility_class,
                "desc": solubility_desc,
            },
            "molecular_properties": {
                "molecular_weight": round(mw, 2),
                "logP": round(logp, 2),
                "tpsa": round(tpsa, 2),
                "hb_donors": hbd,
                "hb_acceptors": hba
            },
            "analysis": {
                "logp_analysis": logp_analysis,
                "mw_analysis": mw_analysis,
                "tpsa_analysis": tpsa_analysis,
                "hydrogen_bonding": f"{hbd}供体/{hba}受体"
            }
        }
        # 构建结果字典
        return result
        
    except Exception as e:
        logging.error(f"计算失败: {str(e)}", exc_info=True)
        return None

def main():
    configure_logging()
    logger = logging.getLogger(__name__)
    
    test_smiles = [
        "C1C(=O)COC2=NC=CN21",  # 有效结构
        "InvalidSMILES",         # 测试错误处理
        "CCO",                   # 乙醇
        "C1=CC=CC=C1C(=O)O"     # 苯甲酸
    ]
    
    for smile in test_smiles:
        try:
            logger.info(f"{'='*40}")
            logger.info(f"Processing SMILES: {smile}")
            
            # 执行预测
            result = predict_water_solubility(smile)
            if not result:
                continue
                
            # 生成可视化
            mol = Chem.MolFromSmiles(smile)
            if mol:
                visualize_molecule(mol)
                
            # 格式化输出
            logger.info(f"预测成功:\n{json.dumps(result, indent=2, ensure_ascii=False)}")
            
        except Exception as e:
            logger.error(f"处理失败: {str(e)}", exc_info=True)

if __name__ == "__main__":
    main()